# memory/enhanced_rag_manager.py
# 增强的Graph RAG管理器，支持三种记忆类型的管理

import asyncio
import yaml
from pathlib import Path
from typing import Optional, Dict, Any, List
from datetime import datetime

from .local_embedding_client import LocalEmbeddingClient, get_local_embedding_client, close_local_embedding_client
from .memory_manager import EnhancedMemoryManager, MemoryType, MemoryImportance
from .rag_retriever import EnhancedRAGRetriever


class EnhancedGraphRAGManager:
    """增强的Graph RAG管理器，负责管理整个增强RAG系统"""
    
    def __init__(self, storage_dir: Path, config: Optional[Dict] = None):
        self.storage_dir = Path(storage_dir)
        
        if config is None:
            # 从统一配置管理器读取
            try:
                from ..configs import get_llm_config
                full_config = get_llm_config()
                config = full_config.get('memory_config', {})
            except ImportError:
                # 回退到直接读取配置文件
                config_path = Path(__file__).parent.parent / "configs" / "llm_config.yaml"
                with open(config_path, 'r', encoding='utf-8') as f:
                    full_config = yaml.safe_load(f)
                    config = full_config.get('memory_config', {})
        
        # 组件
        self.embedding_client: Optional[LocalEmbeddingClient] = None
        self.memory_manager: Optional[EnhancedMemoryManager] = None
        self.rag_retriever: Optional[EnhancedRAGRetriever] = None
        
        # 状态
        self.is_initialized = False
        self.enabled = True  # 是否启用RAG功能
        
        # 存储策略配置
        self.storage_strategy = config.get('storage_strategy', 'selective')  # 'all' 或 'selective'
        self.min_content_length = config.get('min_content_length', 10)  # 最小内容长度
        self.store_short_term = config.get('store_short_term', False)  # 是否存储短期记忆
    
    async def initialize(self) -> bool:
        """
        初始化增强Graph RAG系统
        
        Returns:
            初始化是否成功
        """
        try:
            print("正在初始化增强Graph RAG系统...")
            
            # 1. 初始化本地embedding客户端
            print("初始化本地embedding客户端...")
            print("⏳ 正在加载embedding模型（可能需要一些时间）...")
            try:
                # 注意：模型加载是同步的，可能会阻塞几秒钟
                self.embedding_client = LocalEmbeddingClient()
                print("✅ 本地embedding客户端创建完成")
            except Exception as e:
                print(f"❌ 创建embedding客户端失败: {e}")
                import traceback
                traceback.print_exc()
                self.enabled = False
                return False
            
            # 测试本地embedding连接
            print("测试本地embedding模型...")
            try:
                if not self.embedding_client.test_connection():
                    print("⚠️  警告: 本地embedding模型连接失败，RAG功能将被禁用")
                    self.enabled = False
                    return False
                print("✅ 本地embedding模型测试通过")
            except Exception as e:
                print(f"❌ 测试embedding模型失败: {e}")
                import traceback
                traceback.print_exc()
                self.enabled = False
                return False
            
            # 2. 初始化增强记忆管理器
            try:
                print("初始化增强记忆管理器...")
                self.memory_manager = EnhancedMemoryManager(self.storage_dir, self.embedding_client)
                print("✅ 增强记忆管理器初始化成功")
            except Exception as e:
                print(f"❌ 初始化记忆管理器失败: {e}")
                import traceback
                traceback.print_exc()
                self.enabled = False
                return False
            
            # 3. 初始化增强RAG检索器
            try:
                print("初始化增强RAG检索器...")
                self.rag_retriever = EnhancedRAGRetriever(self.memory_manager, self.embedding_client)
                print("✅ 增强RAG检索器初始化成功")
            except Exception as e:
                print(f"❌ 初始化RAG检索器失败: {e}")
                import traceback
                traceback.print_exc()
                self.enabled = False
                return False
            
            self.is_initialized = True
            print("✅ 增强Graph RAG系统初始化成功")
            
            # 显示embedding客户端信息
            try:
                model_info = self.embedding_client.get_model_info()
                print(f"📊 使用本地embedding模型: {model_info.get('model_path', 'unknown')} (设备: {model_info.get('device', 'unknown')})")
            except Exception as e:
                print(f"⚠️  获取模型信息失败: {e}")
            
            # 显示记忆统计
            try:
                stats = self.memory_manager.get_memory_stats()
                print(f"📚 当前记忆统计: 短期{stats['short_term_count']}个, 长期{stats['long_term_count']}个, 时间{stats['temporal_count']}个")
            except Exception as e:
                print(f"⚠️  获取记忆统计失败: {e}")
            
            return True
            
        except Exception as e:
            print(f"❌ 增强Graph RAG系统初始化失败: {e}")
            import traceback
            traceback.print_exc()
            self.enabled = False
            return False
    
    async def process_conversation(self, user_input: str, ai_response: str) -> bool:
        """
        处理对话，智能决定是否存储到记忆中
        
        Args:
            user_input: 用户输入
            ai_response: AI回复
            
        Returns:
            是否成功添加到记忆
        """
        if not self.is_initialized or not self.enabled:
            return False
        
        try:
            timestamp = datetime.now().isoformat()
            
            # 分析对话内容，决定存储类型和重要性
            user_memory_type, user_importance = self._analyze_content_for_memory(user_input)
            ai_memory_type, ai_importance = self._analyze_content_for_memory(ai_response)
            
            # 批量处理存储
            stored_count = 0
            messages = [
                (user_input, user_memory_type, user_importance, "user_input", "用户输入"),
                (ai_response, ai_memory_type, ai_importance, "ai_response", "AI回复")
            ]
            
            for content, memory_type, importance, tag, description in messages:
                if self._should_store_content(content, memory_type, importance):
                    node_id = await self._add_to_memory(
                        content, 
                        memory_type, 
                        importance,
                        tags=[tag],
                        metadata={"role": tag.split('_')[0], "timestamp": timestamp}
                    )
                    stored_count += 1
                    print(f"{description}已存储到{memory_type.value}记忆: {node_id}")
                else:
                    print(f"{description}未存储（重要性不足）")
            
            if stored_count > 0:
                print(f"对话处理完成，共存储了{stored_count}条记忆")
            else:
                print("对话未存储（内容重要性不足）")
            
            return stored_count > 0
            
        except Exception as e:
            print(f"添加对话到记忆时发生错误: {e}")
            return False
    
    def _analyze_content_for_memory(self, content: str) -> tuple[MemoryType, MemoryImportance]:
        """简化版内容分析 - 只存储有价值的对话"""
        content_lower = content.lower().strip()
        
        # 过滤简单问候和无效内容
        if len(content) < 5 or self._is_simple_question(content_lower):
            return MemoryType.SHORT_TERM, MemoryImportance.LOW
        
        # 事实陈述关键词（用户个人信息）
        factual_keywords = ["我叫", "我的名字", "我喜欢", "我的爱好", "我的职业", "我住在"]
        if any(k in content_lower for k in factual_keywords) and "?" not in content and "？" not in content:
            return MemoryType.LONG_TERM, MemoryImportance.HIGH
        
        # 时间相关内容
        temporal_keywords = ["昨天", "前天", "上周", "之前", "最近", "刚才"]
        if any(k in content_lower for k in temporal_keywords):
            return MemoryType.TEMPORAL, MemoryImportance.MEDIUM
        
        # 默认为短期记忆，中等重要性
        return MemoryType.SHORT_TERM, MemoryImportance.MEDIUM
    
    def _is_simple_question(self, content: str) -> bool:
        """判断是否为简单问候语"""
        simple_patterns = ["你好", "谢谢", "再见", "好的", "嗯", "哦", "啊", "哈哈"]
        content_clean = content.strip().replace("？", "").replace("?", "")
        return content_clean in simple_patterns or len(content_clean) < 3
    
    def _is_question(self, content: str) -> bool:
        """判断是否为疑问句"""
        return "?" in content or "？" in content
    
    
    def _should_store_content(self, content: str, memory_type: MemoryType, importance: MemoryImportance) -> bool:
        """简化版存储判断 - 只存储有意义的内容"""
        content = content.strip()
        
        # 基本检查
        if len(content) < self.min_content_length:
            return False
        
        if self._is_simple_question(content.lower()):
            return False
        
        # 根据记忆类型和重要性决定
        if memory_type == MemoryType.LONG_TERM:
            return importance in [MemoryImportance.HIGH, MemoryImportance.CRITICAL]
        elif memory_type == MemoryType.TEMPORAL:
            return importance >= MemoryImportance.MEDIUM
        else:  # SHORT_TERM
            # 只存储启用短期记忆且重要性足够的内容
            return self.store_short_term and importance >= MemoryImportance.MEDIUM
    
    async def _add_to_memory(self, content: str, memory_type: MemoryType, 
                           importance: MemoryImportance, tags: List[str] = None,
                           metadata: Dict[str, Any] = None) -> str:
        """
        根据类型添加记忆
        
        Args:
            content: 内容
            memory_type: 记忆类型
            importance: 重要性
            tags: 标签
            metadata: 元数据
            
        Returns:
            节点ID
        """
        if memory_type == MemoryType.SHORT_TERM:
            return await self.memory_manager.add_short_term_memory(
                content, tags, importance, metadata
            )
        elif memory_type == MemoryType.LONG_TERM:
            return await self.memory_manager.add_long_term_memory(
                content, tags, importance, metadata
            )
        elif memory_type == MemoryType.TEMPORAL:
            # 提取时间参考
            time_reference = self._extract_time_reference(content)
            return await self.memory_manager.add_temporal_memory(
                content, tags, importance, metadata, time_reference
            )
        else:
            # 默认为短期记忆
            return await self.memory_manager.add_short_term_memory(
                content, tags, importance, metadata
            )
    
    def _extract_time_reference(self, content: str) -> Optional[str]:
        """从内容中提取时间参考"""
        content_lower = content.lower()
        
        time_mappings = {
            "昨天": "yesterday",
            "前天": "day_before_yesterday",
            "今天": "today",
            "上周": "last_week",
            "上个月": "last_month",
            "去年": "last_year",
            "最近": "recent",
            "之前": "before",
            "以前": "before"
        }
        
        for chinese, english in time_mappings.items():
            if chinese in content_lower:
                return english
        
        return None
    
    async def retrieve_context(self, query: str) -> Optional[Dict[str, Any]]:
        """
        检索与查询相关的上下文
        
        Args:
            query: 用户查询
            
        Returns:
            上下文信息或None
        """
        if not self.is_initialized or not self.enabled:
            return None
        
        try:
            # 判断是否应该使用记忆
            should_use = await self.rag_retriever.should_use_memory(query)
            if not should_use:
                return None
            
            # 检索上下文
            context = await self.rag_retriever.retrieve_relevant_context(query)
            
            # 格式化上下文
            formatted_context = await self.rag_retriever.format_context_for_llm(context)
            
            return {
                "context": formatted_context,
                "insights": await self.rag_retriever.get_memory_insights(query),
                "raw_context": context
            }
            
        except Exception as e:
            print(f"检索上下文时发生错误: {e}")
            return None
    
    async def add_user_preference(self, preference: str, category: str = "general") -> str:
        """
        添加用户偏好到长期记忆
        
        Args:
            preference: 偏好内容
            category: 偏好类别
            
        Returns:
            节点ID
        """
        if not self.is_initialized or not self.enabled:
            return ""
        
        try:
            return await self.memory_manager.add_long_term_memory(
                content=preference,
                tags=["user_preference", category],
                importance=MemoryImportance.HIGH,
                metadata={"category": category, "type": "preference"}
            )
        except Exception as e:
            print(f"添加用户偏好时发生错误: {e}")
            return ""
    
    async def add_factual_information(self, fact: str, fact_type: str = "general") -> str:
        """
        添加事实信息到长期记忆
        
        Args:
            fact: 事实内容
            fact_type: 事实类型
            
        Returns:
            节点ID
        """
        if not self.is_initialized or not self.enabled:
            return ""
        
        try:
            return await self.memory_manager.add_long_term_memory(
                content=fact,
                tags=["fact", fact_type],
                importance=MemoryImportance.HIGH,
                metadata={"fact_type": fact_type, "type": "fact"}
            )
        except Exception as e:
            print(f"添加事实信息时发生错误: {e}")
            return ""
    
    async def add_temporal_event(self, event: str, time_reference: str = None) -> str:
        """
        添加时间事件到时间记忆
        
        Args:
            event: 事件内容
            time_reference: 时间参考
            
        Returns:
            节点ID
        """
        if not self.is_initialized or not self.enabled:
            return ""
        
        try:
            return await self.memory_manager.add_temporal_memory(
                content=event,
                tags=["event", "temporal"],
                importance=MemoryImportance.MEDIUM,
                metadata={"type": "event"},
                time_reference=time_reference
            )
        except Exception as e:
            print(f"添加时间事件时发生错误: {e}")
            return ""
    
    async def get_memory_stats(self) -> Dict[str, Any]:
        """获取记忆统计信息"""
        if not self.is_initialized or not self.enabled:
            return {"enabled": False}
        
        stats = self.memory_manager.get_memory_stats()
        stats["enabled"] = True
        stats["initialized"] = self.is_initialized
        return stats
    
    async def clear_memory(self, memory_type: MemoryType = None):
        """
        清空记忆
        
        Args:
            memory_type: 要清空的记忆类型，None表示清空所有
        """
        if self.memory_manager:
            self.memory_manager.clear_memory(memory_type)
            if memory_type is None:
                print("所有记忆已清空")
            else:
                print(f"{memory_type.value}记忆已清空")
    
    async def search_memories(self, query: str, memory_types: List[MemoryType] = None, 
                            limit: int = 10) -> List[tuple]:
        """
        搜索记忆
        
        Args:
            query: 查询文本
            memory_types: 要搜索的记忆类型列表
            limit: 返回结果数量限制
            
        Returns:
            记忆搜索结果列表
        """
        if not self.is_initialized or not self.enabled:
            return []
        
        try:
            return await self.memory_manager.search_memories(query, memory_types, limit)
        except Exception as e:
            print(f"搜索记忆时发生错误: {e}")
            return []
    
    async def close(self):
        """关闭增强Graph RAG系统 - 添加异常保护"""
        try:
            if self.embedding_client:
                close_local_embedding_client()
            print("✓ Graph RAG系统已关闭")
        except Exception as e:
            print(f"关闭RAG系统错误: {e}")


# 全局增强Graph RAG管理器实例
_enhanced_rag_manager: Optional[EnhancedGraphRAGManager] = None


async def initialize_enhanced_graph_rag(storage_dir: Path, config: Optional[Dict] = None) -> EnhancedGraphRAGManager:
    """
    初始化全局增强Graph RAG管理器
    
    Args:
        storage_dir: 存储目录
        config: 配置字典，如果为None则从配置文件读取
        
    Returns:
        增强Graph RAG管理器实例
    """
    global _enhanced_rag_manager
    if _enhanced_rag_manager is None:
        _enhanced_rag_manager = EnhancedGraphRAGManager(storage_dir, config)
        await _enhanced_rag_manager.initialize()
    return _enhanced_rag_manager


async def get_enhanced_rag_manager() -> Optional[EnhancedGraphRAGManager]:
    """获取全局增强Graph RAG管理器实例"""
    return _enhanced_rag_manager


async def close_enhanced_graph_rag():
    """关闭全局增强Graph RAG管理器"""
    global _enhanced_rag_manager
    if _enhanced_rag_manager:
        await _enhanced_rag_manager.close()
        _enhanced_rag_manager = None
